package com.ibm.risk.irmp.common.utils;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * SSH远程执行命令
 * 仅运行执行限定的查看、监控等命令
 * 不允许执行shell脚本、动态脚本、删除命令等可能对系统造成危害以及修改文件和数据的命令
 */
public class SshUtil {
    private final static Logger log = LoggerFactory.getLogger(SshUtil.class);
    private final static Integer TIMEOUT_SECONDS_MAX = 60;
    private final static Integer TIMEOUT_SECONDS = 10;

    //允许执行的命令
    private final static List<String> allowedCommandList = Arrays.asList("cd", "ls", "ps", "echo", "tail", "more", "less",
            "cat", "head", "grep", "jps", "whereis", "which", "find", "free", "sleep", "wc", "df", "du",
            "tar", "jar", "zip", "unzip",
            "hdfs dfs -ls", "hdfs dfs -get", "hdfs dfs -mkdir", "hdfs dfs -chmod", "hdfs dfs -chown", "hdfs dfs -chgrp",
            "hdfs dfs -put", "hdfs dfs -cp", "hdfs dfs -tail",
            "dos2unix ", "unix2dos", "cp", "mv", "mkdir"
            , "kkkill"
    );


    public static String exec(String host, String username, String password, String command) {
        return exec(host, username, password, command, TIMEOUT_SECONDS_MAX, TimeUnit.SECONDS);
    }

    public static String exec(String host, String username, String password, String command, long timeout, TimeUnit unit) {
        String result = "";
        BasicThreadFactory build = new BasicThreadFactory.Builder()
                .daemon(false)
                .namingPattern("exec-command")
                .build();
        ExecutorService executorService = Executors.newSingleThreadExecutor(build);
        Future<String> submit = executorService.submit(new Callable<String>() {
            @Override
            public String call() throws Exception {
                return execCommand(host, username, password, command);
            }
        });
        try {
            result = submit.get(timeout, unit);
            executorService.shutdown();
            //} catch (TimeoutException | InterruptedException | ExecutionException exp) {
        } catch (Exception exp) {
            //exp.printStackTrace();
            // 超时的时候向线程池中所有的线程发出中断(interrupted)。
            executorService.shutdownNow();
            throw new RuntimeException(exp.getCause() != null ? exp.getCause().getMessage() : exp.getLocalizedMessage(), exp);
        } finally {
            // Harmless if task already completed
            submit.cancel(true); // interrupt if running
        }
        if (result == null || result.length() == 0) {
            result += "执行成功 ";
        }
        return result;
    }

    private static String execCommand(String host, String username, String password, String command) {
        log.info("host: {}, username: {},  exec command:  {}", host, username, command);
        if (command == null || command.trim().length() == 0) {
            throw new RuntimeException("处理命令不能为空");
        }
        if (validate(command) == false) {
            //throw new RuntimeException("仅允许执行限定的命令(" + allowedCommandList.toString() + ")");
            throw new RuntimeException("仅允许执行限定的命令");
        }
        if (StringUtils.startsWith(command, "kkkill")) {
            command = StringUtils.replace(command, "kkkill", "kill");
        }

        // 2>&1
        String execCommand = command + " 2>&1 & ";
        //获取进程号
        //String pidCommand = "pid=$(ps aux | grep \"" + command + "\" | awk '{print $2}' | sort -n | head -n 1) ";
        //延迟180秒后执行kill命令，关闭ssh进程，延迟时间可以根据调用的命令不同调整
        //kill -9 `ps -ef |grep tail | awk '{print $2}'`
        //execCommand = execCommand + "sleep 10 && kill -9 `ps -ef |grep " + command + " | awk '{print $2}'`";
        log.info("execCommand: {}", command);

        StringBuilder info = new StringBuilder();
        Session session = null;
//        ChannelShell channel = null;
        ChannelExec channel = null;
        InputStream in = null;
        BufferedReader br = null;
        try {
            Properties properties = new Properties();
            //第一次访问服务器不用输入yes
            properties.put("StrictHostKeyChecking", "no");

            JSch jsch = new JSch();
            session = jsch.getSession(username, host);
            session.setPassword(password);
            session.setConfig(properties);
            session.setTimeout(TIMEOUT_SECONDS * 1000);
            session.connect();

//            channel = (ChannelShell) session.openChannel("shell");
            channel = (ChannelExec) session.openChannel("exec");
            channel.setPty(false);
            //channel.setCommand(execCommand);
            //JSCH exec方式加载环境变量
            channel.setCommand("source /etc/profile > /dev/null && source ~/.bash_profile > /dev/null && source ~/.bashrc > /dev/null &&  " + execCommand);
            channel.setErrStream(System.err);
            channel.connect();

            //后台运行的最后一个进程的ID号
            //OutputStream outputStream = channel.getOutputStream();
            //String grepCommand = "kill -9 `ps -ef |grep " + command + " | awk '{print $2}'`";
//            String grepCommand = "pid=$! && kill -9 ${pid}";
            //outputStream.write(grepCommand.getBytes());
//            channel.setCommand(grepCommand);

            in = channel.getInputStream();
            InputStreamReader inputStr = new InputStreamReader(in, StandardCharsets.UTF_8);
            br = new BufferedReader(inputStr);

            long startTime = (new Date()).getTime();
            String line = "";
            while ((line = br.readLine()) != null) {
                info.append(line).append("\r\n");
                //System.out.println(line);
                //最长n秒后返回，但是如果没有数据此处判断无法触发
                if ((new Date()).getTime() - startTime > TIMEOUT_SECONDS * 1000) {
                    break;
                }
            }
        } catch (JSchException exp) {
            exp.printStackTrace();
            throw new RuntimeException("连接失败", exp);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
            if (channel != null && !channel.isClosed()) {
                channel.disconnect();
            }
            if (in != null) {
                try {
                    in.close();
                } catch (IOException e) {
                }
            }
            if (br != null) {
                try {
                    br.close();
                } catch (IOException e) {
                }
            }
        }
        return info.toString().trim();
    }

    /**
     * 验证命令，限定在指定的范围之内
     *
     * @param command
     * @return
     */
    public static boolean validate(String command) {
        if (command == null || "".equals(command.trim())) {
            return false;
        }

        //分割命令 |  ;  &&
        List<String> cmds = new ArrayList<>();
        Arrays.stream(command.split("\\|")).forEachOrdered(x1 -> {
            Arrays.stream(x1.split(";")).forEachOrdered(x2 -> {
                Arrays.stream(x2.split("&&")).forEachOrdered(x3 -> {
                    if (x3 != null && !"".equals(x3.trim())) {
                        cmds.add(x3.trim() + " ");
                    }
                });
            });
        });

        AtomicBoolean isAllowd = new AtomicBoolean(true);
        cmds.forEach(cmd -> {
            AtomicBoolean isAllowd2 = new AtomicBoolean(false);
            allowedCommandList.forEach(x -> {
                if (cmd.startsWith(x + " ")) {
                    isAllowd2.set(true);
                }
            });
            if (isAllowd2.get() == false) {
                isAllowd.set(false);
            }
        });
        if (isAllowd.get() == false) {
            return false;
        }

        //禁止执行 tailf | tail -f | rm | `间隔号 | $(...)
        cmds.forEach(cmd -> {
//            if(cmd.startsWith("tail ")) {
//                if(cmd.indexOf("-f ") > 0) {
//                    isAllowd.set(false);
//                }
//            }
//            if(cmd.startsWith("tailf ") ) {
//                isAllowd.set(false);
//            }
            if (cmd.startsWith("rm ")) {
                isAllowd.set(false);
            }
            if (cmd.indexOf("`") >= 0) {
                isAllowd.set(false);
            }
            if (cmd.indexOf("$(") >= 0) {
                isAllowd.set(false);
            }
        });
        if (isAllowd.get() == false) {
            return false;
        }

        return true;
    }
}

